#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
标注同步模块
负责将单个 Labelme JSON 文件同步到 COCO 格式文件
"""

import json
import os
import shutil
import tempfile
from pathlib import Path
from typing import List, Optional, Tuple
from PyQt5.QtCore import QPointF


def sync_all_json_to_coco(annotations_dir: Path) -> Tuple[int, int]:
    """
    同步所有单个 JSON 文件到对应的 COCO 文件
    
    Args:
        annotations_dir: 标注文件目录
        
    Returns:
        (成功数量, 失败数量)
    """
    if not annotations_dir.exists():
        return 0, 0
    
    # 查找所有单个JSON文件
    json_files = list(annotations_dir.glob('*.json'))
    # 排除COCO文件
    json_files = [f for f in json_files if f.name not in [
        'instance_train.json', 'instance_val.json', 
        'instance_all.json', 'annotations.json'
    ]]
    
    if not json_files:
        return 0, 0
    
    print(f"开始同步 {len(json_files)} 个单个JSON文件到COCO...")
    
    # 查找COCO文件
    coco_files = [
        annotations_dir / 'instance_train.json',
        annotations_dir / 'instance_val.json',
        annotations_dir / 'instance_all.json',
    ]
    
    success_count = 0
    fail_count = 0
    
    for json_file in json_files:
        try:
            # 从JSON文件名推断图片名
            image_name = json_file.stem + '.jpg'  # 假设是jpg，需要从JSON中获取实际名称
            
            # 读取JSON文件获取图片名
            try:
                with open(json_file, 'r', encoding='utf-8') as f:
                    label_data = json.load(f)
                image_name = label_data.get('imagePath', image_name)
                # 处理相对路径，提取文件名
                if '/' in image_name or '\\' in image_name:
                    image_name = Path(image_name).name
            except Exception as e:
                print(f"⚠ 读取 {json_file.name} 失败: {e}")
                fail_count += 1
                continue
            
            # 找到对应的COCO文件
            coco_file = _find_coco_file_for_image(image_name, coco_files, annotations_dir)
            
            if not coco_file:
                print(f"⚠ 未找到 {image_name} 对应的COCO文件，跳过")
                fail_count += 1
                continue
            
            # 同步单个JSON到COCO
            if sync_single_json_to_coco(json_file, coco_file, annotations_dir):
                success_count += 1
            else:
                fail_count += 1
                
        except Exception as e:
            print(f"⚠ 同步 {json_file.name} 失败: {e}")
            fail_count += 1
    
    if success_count > 0:
        print(f"✓ 已同步 {success_count} 个JSON文件到COCO")
    if fail_count > 0:
        print(f"⚠ {fail_count} 个JSON文件同步失败")
    
    return success_count, fail_count


def _find_coco_file_for_image(image_name: str, coco_files: List[Path], annotations_dir: Path) -> Optional[Path]:
    """
    查找包含指定图片的 COCO 文件
    
    Args:
        image_name: 图片文件名
        coco_files: COCO 文件列表
        annotations_dir: 标注目录（用于默认值）
        
    Returns:
        COCO 文件路径，如果未找到则返回 None
    """
    image_name_normalized = Path(image_name).name
    
    for cf in coco_files:
        if not cf.exists():
            continue
        
        try:
            with open(cf, 'r', encoding='utf-8') as f:
                coco_data = json.load(f)
            
            for img in coco_data.get('images', []):
                # 规范化文件名匹配
                coco_file_name = img.get('file_name', '')
                coco_file_name_normalized = Path(coco_file_name).name
                if coco_file_name_normalized == image_name_normalized or coco_file_name == image_name:
                    return cf
        except Exception as e:
            print(f"⚠ 读取 COCO 文件 {cf.name} 失败: {e}")
            continue
    
    # 如果没有找到，默认使用 instance_train.json（如果存在）
    default_coco = annotations_dir / 'instance_train.json'
    if default_coco.exists():
        return default_coco
    
    return None


def sync_single_json_to_coco(json_file: Path, coco_file: Path, annotations_dir: Path) -> bool:
    """
    同步单个 JSON 文件到 COCO 文件
    
    Args:
        json_file: Labelme 格式的 JSON 文件路径
        coco_file: COCO 格式的 JSON 文件路径
        annotations_dir: 标注目录（用于构建图片路径）
        
    Returns:
        是否成功
    """
    try:
        # 读取 JSON 文件
        with open(json_file, 'r', encoding='utf-8') as f:
            label_data = json.load(f)
        
        if not label_data or 'shapes' not in label_data:
            print(f"⚠ {json_file.name} 没有有效的标注数据")
            return False
        
        # 获取图片名
        image_name = label_data.get('imagePath', json_file.stem + '.jpg')
        if '/' in image_name or '\\' in image_name:
            image_name = Path(image_name).name
        
        # 构建图片路径
        image_path = annotations_dir.parent / 'images' / image_name
        
        # 转换 shapes
        shapes = []
        for shape_dict in label_data.get('shapes', []):
            try:
                # 创建临时Shape对象（用于兼容 _update_coco_file 函数）
                from labelme.app import Shape
                shape = Shape(
                    label=shape_dict.get('label', ''),
                    shape_type=shape_dict.get('shape_type', 'rectangle'),
                    flags=shape_dict.get('flags', {})
                )
                # 转换points
                points = shape_dict.get('points', [])
                shape.points = [QPointF(float(p[0]), float(p[1])) for p in points if len(p) >= 2]
                shapes.append(shape)
            except Exception as e:
                print(f"⚠ 转换 shape 失败: {e}")
                continue
        
        if not shapes:
            print(f"⚠ {json_file.name} 没有有效的 shapes")
            return False
        
        # 更新COCO文件
        image_width = label_data.get('imageWidth', 0)
        image_height = label_data.get('imageHeight', 0)
        
        # 导入 _update_coco_file 函数（从 training_plugin 中）
        # 为了避免循环导入，我们在这里直接实现更新逻辑
        return _update_coco_file_direct(coco_file, image_path, shapes, image_width, image_height)
        
    except Exception as e:
        print(f"⚠ 同步 {json_file.name} 到 COCO 失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def _update_coco_file_direct(coco_file: Path, image_path: Path, shapes: List, image_width: int, image_height: int) -> bool:
    """
    直接更新 COCO 文件（不依赖 training_plugin）
    
    Args:
        coco_file: COCO 文件路径
        image_path: 图片文件路径
        shapes: Labelme Shape 对象列表
        image_width: 图片宽度
        image_height: 图片高度
        
    Returns:
        是否成功
    """
    try:
        # 加载现有 COCO 数据
        if coco_file.exists():
            with open(coco_file, 'r', encoding='utf-8') as f:
                coco_data = json.load(f)
        else:
            # 创建新的 COCO 数据结构
            coco_data = {
                'info': {
                    'description': 'Auto Generated Dataset',
                    'version': '1.0',
                    'year': 2024
                },
                'licenses': [],
                'images': [],
                'annotations': [],
                'categories': []
            }
        
        # 获取或创建类别映射
        categories = {}
        category_id_map = {}
        for cat in coco_data.get('categories', []):
            categories[cat['name']] = cat['id']
            category_id_map[cat['id']] = cat['name']
        
        # 查找或创建图片记录
        image_name = image_path.name
        image_info = None
        for img in coco_data.get('images', []):
            if Path(img.get('file_name', '')).name == image_name:
                image_info = img
                break
        
        if not image_info:
            # 创建新图片记录
            next_image_id = max([img['id'] for img in coco_data.get('images', [])], default=0) + 1
            image_info = {
                'id': next_image_id,
                'file_name': image_name,
                'width': image_width,
                'height': image_height
            }
            coco_data['images'].append(image_info)
        else:
            # 更新现有图片信息
            image_info['width'] = image_width
            image_info['height'] = image_height
        
        image_id = image_info['id']
        
        # 删除该图片的旧标注
        coco_data['annotations'] = [
            ann for ann in coco_data.get('annotations', [])
            if ann.get('image_id') != image_id
        ]
        
        # 添加新标注
        next_annotation_id = max([ann['id'] for ann in coco_data.get('annotations', [])], default=0) + 1
        
        for shape in shapes:
            # 获取类别
            if hasattr(shape, 'label'):
                label = shape.label
            else:
                label = shape.get('label', 'unknown')
            
            # 获取或创建类别ID
            if label not in categories:
                next_category_id = max(categories.values(), default=0) + 1
                categories[label] = next_category_id
                category_id_map[next_category_id] = label
                coco_data['categories'].append({
                    'id': next_category_id,
                    'name': label,
                    'supercategory': 'none'
                })
            
            category_id = categories[label]
            
            # 转换标注
            annotation = _convert_shape_to_coco_annotation(
                shape, image_id, next_annotation_id, category_id, image_width, image_height
            )
            
            if annotation:
                coco_data['annotations'].append(annotation)
                next_annotation_id += 1
        
        # 更新日期
        from datetime import datetime
        coco_data['info']['date_created'] = datetime.now().strftime('%Y-%m-%d')
        
        # 安全保存文件
        temp_file = coco_file.with_suffix('.tmp')
        try:
            with open(temp_file, 'w', encoding='utf-8') as f:
                json.dump(coco_data, f, ensure_ascii=False, indent=2)
                f.flush()
                os.fsync(f.fileno())
            
            shutil.move(str(temp_file), str(coco_file))
            return True
        except Exception as e:
            if temp_file.exists():
                try:
                    temp_file.unlink()
                except:
                    pass
            raise
        
    except Exception as e:
        print(f"⚠ 更新 COCO 文件失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def _convert_shape_to_coco_annotation(shape, image_id: int, annotation_id: int, category_id: int, 
                                      image_width: int, image_height: int) -> Optional[dict]:
    """
    将 Labelme Shape 转换为 COCO 格式标注
    
    Args:
        shape: Labelme Shape 对象
        image_id: 图片ID
        annotation_id: 标注ID
        category_id: 类别ID
        image_width: 图片宽度
        image_height: 图片高度
        
    Returns:
        COCO 格式的标注字典
    """
    try:
        # 获取 points
        if hasattr(shape, 'points'):
            points = shape.points
        else:
            points = shape.get('points', [])
        
        if not points or len(points) < 2:
            return None
        
        # 转换为坐标列表
        coords = []
        for point in points:
            if hasattr(point, 'x') and hasattr(point, 'y'):
                coords.append([float(point.x()), float(point.y())])
            elif isinstance(point, (list, tuple)) and len(point) >= 2:
                coords.append([float(point[0]), float(point[1])])
        
        if len(coords) < 2:
            return None
        
        # 获取 shape_type
        if hasattr(shape, 'shape_type'):
            shape_type = shape.shape_type
        else:
            shape_type = shape.get('shape_type', 'rectangle')
        
        # 转换为 COCO 格式
        if shape_type == 'rectangle' and len(coords) == 2:
            # 矩形框：[x, y, width, height]
            x1, y1 = coords[0]
            x2, y2 = coords[1]
            x = min(x1, x2)
            y = min(y1, y2)
            w = abs(x2 - x1)
            h = abs(y2 - y1)
            
            bbox = [x, y, w, h]
            area = w * h
            
            # 从bbox生成顺时针的segmentation（矩形四个角点）
            # 顺时针顺序：左上 -> 右上 -> 右下 -> 左下
            segmentation = [[
                float(x),           # 左上 x
                float(y),           # 左上 y
                float(x + w),       # 右上 x
                float(y),           # 右上 y
                float(x + w),       # 右下 x
                float(y + h),       # 右下 y
                float(x),           # 左下 x
                float(y + h)        # 左下 y
            ]]
        else:
            # 多边形或其他形状
            xs = [c[0] for c in coords]
            ys = [c[1] for c in coords]
            x = min(xs)
            y = min(ys)
            w = max(xs) - x
            h = max(ys) - y
            
            bbox = [x, y, w, h]
            # 计算多边形面积（简化：使用边界框面积）
            area = w * h
            
            # 对于多边形，将坐标转换为segmentation格式 [x1, y1, x2, y2, ...]
            segmentation = [[float(coord) for point in coords for coord in point]]
            
            # 如果多边形坐标少于3个点，使用bbox生成矩形segmentation
            if len(coords) < 3:
                segmentation = [[
                    float(x),           # 左上 x
                    float(y),           # 左上 y
                    float(x + w),       # 右上 x
                    float(y),           # 右上 y
                    float(x + w),       # 右下 x
                    float(y + h),       # 右下 y
                    float(x),           # 左下 x
                    float(y + h)        # 左下 y
                ]]
        
        return {
            'id': annotation_id,
            'image_id': image_id,
            'category_id': category_id,
            'bbox': bbox,
            'area': area,
            'iscrowd': 0,
            'segmentation': segmentation
        }
        
    except Exception as e:
        print(f"⚠ 转换标注失败: {e}")
        return None

